import json
import time
import threading
import queue
import logging
from hqyj_mqtt import Mqtt_Clt
import base64
import numpy as np
import cv2
from predict import predict_yolo

# 配置日志
# 设置日志记录的基本配置，包括日志级别、格式和输出目标
# 日志级别设置为INFO，格式包含时间戳、日志级别和消息内容
# 日志同时输出到文件'cargo_sorting.log'和控制台
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.FileHandler('cargo_sorting.log'),
        logging.StreamHandler()
    ]
)
# 创建日志记录器实例，用于记录程序运行过程中的信息
logger = logging.getLogger(__name__)


# 推杆与传感器定义
# 定义推杆和传感器的常量标识符
load_rod = "first"          # 装载推杆标识
ripe_rod = "second"         # 成熟推杆标识
half_ripe_rod = "third"     # 半成熟推杆标识
raw_rod = "fourth"          # 生推杆标识
FIRST_SWITCH = "first_switch"     # 第一传感器开关标识
SECOND_SWITCH = "second_switch"   # 第二传感器开关标识
THIRD_SWITCH = "third_switch"     # 第三传感器开关标识


# 传感器顺序（基于流水线从左到右）
SENSOR_ORDER = ["first_switch", "second_switch", "third_switch"] # 传感器顺序

# 分类-传感器映射
CLASS_SWITCH_MAP = {
    "ripe": FIRST_SWITCH,
    "half-ripe": SECOND_SWITCH,
    "raw": THIRD_SWITCH
}

# 分类-推杆映射
CLASS_ROD_MAP = {
    "ripe": ripe_rod,
    "half-ripe": half_ripe_rod,
    "raw": raw_rod
}

# 配置参数
PUSH_DELAY = 0.2  # 推杆延迟时间
RECOGNIZE_INTERVAL = 2  # 上货间隔
ROD_ACTION_TIME = { # 推杆动作时间
    "first": 0.6,
    "second": 0.6,
    "third": 0.6,
    "fourth": 0.6
}
SENSOR_TIMEOUT = 8.0  # 传感器超时时间
CARGO_TIMEOUT = 30.0  # 货物超时清理时间

# 创建两个队列用于存储杆件数据
# load_rod_queue: 用于存储需要处理的杆件数据，最大容量为5个元素
load_rod_queue = queue.Queue(maxsize=5)
# other_rod_queue: 用于存储其他杆件数据，最大容量为10个元素
other_rod_queue = queue.Queue(maxsize=10)


# 上货锁
# 全局变量定义部分
# load_lock: 线程锁，用于保护加载操作的线程安全
# load_timer: 定时器对象，用于控制加载操作的时间间隔
# last_load_time: 上次加载操作的时间戳，用于避免频繁加载
# last_sensor_trigger_time: 字典类型，以传感器ID为键，记录每个传感器最后一次触发的时间戳

load_lock = threading.Lock()
load_timer = None
last_load_time = 0.0
last_sensor_trigger_time = {}  # 记录每个传感器最后触发时间



class CargoTracker:
    """
    货物追踪器类

    用于跟踪和管理货物的状态信息，包括分类、触发开关、状态等属性
    """

    def __init__(self):
        """
        初始化货物追踪器

        初始化一个空的货物映射表和货物ID计数器
        """
        self.cargo_map = {}  # {cargo_id: {classification, triggered_switch, status, create_time, sensor_triggered}}
        self.next_cargo_id = 1


    def create_new_cargo(self, classification):
        """
        创建一个新的货物记录

        :param classification: 货物分类标识
        :return: 新创建货物的ID
        """
        # 生成唯一的货物ID并更新下一个可用ID
        cargo_id = self.next_cargo_id
        self.next_cargo_id += 1

        # 初始化货物信息并存储到映射表中
        self.cargo_map[cargo_id] = {
            "classification": classification,
            "triggered_switch": None,
            "status": "pending",
            "create_time": time.time(),
            "sensor_triggered": False
        }

        # 记录新货物创建日志
        logger.info(f"生成新货物ID: {cargo_id}, 分类: {classification}, 应触发传感器: {CLASS_SWITCH_MAP[classification]}")
        return cargo_id


    def update_cargo_switch(self, switch_name):
        """
        更新货物开关状态，包含防抖处理

        参数:
            switch_name (str): 开关名称标识符

        返回值:
            None: 当触发被忽略时返回None
        """
        current_time = time.time()
        # 防抖处理：忽略310毫秒内的重复触发
        if switch_name in last_sensor_trigger_time and (current_time - last_sensor_trigger_time[switch_name] < 0.2): # 当短时间内多次触发同一任务时，只保留最后一次触发的任务
            logger.warning(f"传感器 {switch_name} 触发过于频繁，忽略")
            return None
        last_sensor_trigger_time[switch_name] = current_time


        # 遍历所有货物，查找与当前开关匹配的待处理货物
        # 当货物状态为等待处理、分类匹配且未被传感器触发时，标记该货物并启动超时重置定时器
        for cargo_id, cargo_info in list(self.cargo_map.items()):
            # 检查货物是否满足触发条件：状态为pending、分类匹配、且未被传感器触发
            if (cargo_info["status"] == "pending" and
                CLASS_SWITCH_MAP.get(cargo_info["classification"]) == switch_name and
                not cargo_info["sensor_triggered"]):
                # 标记触发的开关名称
                cargo_info["triggered_switch"] = switch_name
                # 更新货物状态为处理中
                cargo_info["status"] = "processing"
                # 标记传感器已触发
                cargo_info["sensor_triggered"] = True
                # 启动超时重置定时器，在SENSOR_TIMEOUT秒后执行reset_sensor方法
                threading.Timer(SENSOR_TIMEOUT, self.reset_sensor, args=(switch_name, cargo_id)).start()
                # 记录货物触发传感器的日志信息
                logger.info(f"货物ID: {cargo_id}（{cargo_info['classification']}）触发匹配传感器: {switch_name}")
                # 返回已触发的货物ID
                return cargo_id
        # 如果没有找到匹配的货物，记录警告日志并返回None
        logger.warning(f"无匹配分类的待处理货物，忽略传感器: {switch_name}")
        return None


    def reset_sensor(self, switch_name, cargo_id):
        """
        重置指定货物的传感器触发状态

        :param switch_name: 开关名称，用于匹配触发的开关
        :param cargo_id: 货物ID，用于定位需要重置的货物
        :return: 无返回值
        """
        # 检查货物是否存在且触发开关匹配时，重置传感器触发状态
        if cargo_id in self.cargo_map and self.cargo_map[cargo_id]["triggered_switch"] == switch_name:
            self.cargo_map[cargo_id]["sensor_triggered"] = False
            logger.info(f"货物ID: {cargo_id} 的传感器 {switch_name} 已重置")

    def mark_cargo_done(self, cargo_id):
        """
        标记货物处理完成

        参数:
            cargo_id: 货物ID，用于标识需要标记完成的货物

        返回值:
            无返回值
        """
        # 检查货物ID是否存在于货物映射表中
        if cargo_id in self.cargo_map:
            # 获取触发的开关名称
            switch_name = self.cargo_map[cargo_id]["triggered_switch"]
            # 如果存在触发的开关，则重置传感器状态
            if switch_name:
                self.reset_sensor(switch_name, cargo_id)
            # 记录货物处理完成的日志信息
            logger.info(f"货物ID: {cargo_id} 处理完成")
            # 从货物映射表中删除该货物记录
            del self.cargo_map[cargo_id]


    def cleanup_old_cargos(self):
        """
        清理超时的货物信息

        该函数会遍历cargo_map中的所有货物，检查每个货物的创建时间，
        如果超过预设的超时时间(CARGO_TIMEOUT)，则将该货物从映射表中移除。

        参数:
            无

        返回值:
            无
        """
        current_time = time.time()
        to_remove = []
        # 遍历所有货物，找出超时的货物ID
        for cargo_id, cargo_info in self.cargo_map.items():
            if current_time - cargo_info["create_time"] > CARGO_TIMEOUT:
                logger.warning(f"货物ID: {cargo_id} 超时未处理，移除")
                to_remove.append(cargo_id)
        # 批量删除超时的货物信息
        for cargo_id in to_remove:
            del self.cargo_map[cargo_id]


def safe_send_json(mqtt_client, msg):
    """
    安全发送JSON消息到MQTT服务器

    参数:
        mqtt_client: MQTT客户端对象，需要实现send_json_msg和reconnect方法
        msg: 要发送的JSON消息内容

    返回值:
        无返回值
    """
    try:
        # 尝试发送JSON消息
        mqtt_client.send_json_msg(msg)
    except Exception as e:
        # 记录发送失败的错误信息
        logger.error(f"MQTT 发送失败: {e}")
        try:
            # 发送失败后尝试重新连接
            mqtt_client.reconnect()
        except:
            # 记录重连失败的错误信息
            logger.error("MQTT 重连失败")


def push_pull(rod, mqtt_client, delay=0):
    """
    控制推杆执行推拉动作

    参数:
        rod: 推杆标识符，用于指定要控制的推杆
        mqtt_client: MQTT客户端对象，用于发送控制命令
        delay: 延迟执行时间（秒），默认为0表示立即执行

    返回值:
        无返回值

    功能说明:
        该函数首先根据指定的延迟时间等待，然后通过MQTT客户端发送推杆的推出命令，
        等待推杆动作完成后再发送拉回命令，实现推杆的完整推拉动作控制
    """
    # 如果设置了延迟时间，则等待指定秒数后再执行推杆动作
    # 如果延迟时间大于0，则记录等待信息并执行延时
    if delay > 0:
        logger.info(f"等待 {delay:.2f}秒 后执行推杆 {rod}")
        time.sleep(delay)
    # 记录执行推杆动作的信息
    logger.info(f"执行推杆 {rod} 动作")

    try:
        # 发送推杆推出命令，并等待动作完成后再发送拉回命令
         # 发送推杆推出命令
        safe_send_json(mqtt_client, json.dumps({"rod_control": f"{rod}_push"}))
        logger.info(f"发送 {rod}_push 命令")
        # 等待推杆动作完成
        time.sleep(ROD_ACTION_TIME[rod])
        # 发送推杆收回命令
        safe_send_json(mqtt_client, json.dumps({"rod_control": f"{rod}_pull"}))
        logger.info(f"发送 {rod}_pull 命令")
    # 异常处理块：捕获推杆命令执行过程中的所有异常
    # 当推杆命令执行失败时，记录错误日志，包含推杆标识和具体异常信息
    except Exception as e:
        logger.error(f"推杆 {rod} 命令失败: {e}")



def load_rod_worker(mqtt_client):
    """
    上货推杆工作线程函数，持续从队列中获取上货任务并执行

    参数:
        mqtt_client: MQTT客户端实例，用于与MQTT服务器通信

    返回值:
        无返回值
    """
    while True:
        try:
            # 从上货队列中获取任务，超时时间为5秒
            rod = load_rod_queue.get(timeout=5.0)
            # 如果获取到None标识，则退出工作循环
            if rod is None:
                break
            # 执行推杆动作
            push_pull(rod, mqtt_client)
            # 标记当前任务完成
            load_rod_queue.task_done()
            # 记录任务完成日志
            logger.info(f"上货推杆 {rod} 任务完成，当前队列大小: {load_rod_queue.qsize()}")
        except queue.Empty:
            # 队列为空时继续循环等待
            continue
        except Exception as e:
            # 记录其他异常错误
            logger.error(f"load_rod_worker 错误: {e}")


def other_rod_worker(mqtt_client, tracker):
    """
    其他推杆工作线程函数，用于处理分类推杆任务队列中的任务

    该函数作为一个后台工作线程持续运行，从other_rod_queue队列中获取分类推杆任务，
    执行推杆操作并通过MQTT客户端发送控制指令，完成后更新货物跟踪状态。

    参数:
        mqtt_client: MQTT客户端对象，用于发送推杆控制指令
        tracker: 货物跟踪器对象，用于标记货物处理完成状态

    返回值:
        无返回值
    """
    while True:
        try:
            # 从分类推杆任务队列中获取任务，超时时间为5秒
            task = other_rod_queue.get(timeout=5.0)
            if task is None:
                break
            cargo_id, rod, switch_name = task

            # 执行推杆推送操作
            push_pull(rod, mqtt_client, delay=PUSH_DELAY)

            # 标记货物处理完成并更新队列状态
            tracker.mark_cargo_done(cargo_id)
            other_rod_queue.task_done()

            # 记录任务完成日志
            logger.info(f"分类推杆 {rod} 任务完成，当前队列大小: {other_rod_queue.qsize()}")
        except queue.Empty:
            # 队列为空时继续等待新任务
            continue
        except Exception as e:
            # 记录处理过程中的异常错误
            logger.error(f"other_rod_worker 错误: {e}")


def start_workers(mqtt_client, tracker):
    # 创建工作线程列表和负载数据加载线程
    # 启动一个后台线程来处理ROD负载数据的加载工作，并将该线程添加到工作线程列表中
    workers = []
    load_thread = threading.Thread(target=load_rod_worker, args=(mqtt_client,), daemon=True)
    load_thread.start()
    workers.append(load_thread)

    # 创建并启动两个工作线程来处理其他杆设备的任务
    # 每个线程都以守护线程模式运行，执行other_rod_worker函数
    for _ in range(2):
        t = threading.Thread(target=other_rod_worker, args=(mqtt_client, tracker), daemon=True)
        t.start()
        workers.append(t)
    return workers


def trigger_load(mqtt_client, retry_count=0, max_retries=3):
    # 获取全局变量load_timer和last_load_time的引用
    global load_timer, last_load_time
    # 使用加载锁确保线程安全，重置加载计时器并获取当前时间戳
    with load_lock:
        load_timer = None
        current_time = time.time()

        # 检查上货间隔时间是否满足要求，避免频繁触发上货操作
        # 如果距离上次上货时间小于设定的识别间隔，且不是第一次上货，则需要等待一段时间后再触发
        if current_time - last_load_time < RECOGNIZE_INTERVAL and last_load_time != 0:
            # 计算需要等待的时间
            wait_time = RECOGNIZE_INTERVAL - (current_time - last_load_time)
            logger.info(f"上货间隔未达标，等待 {wait_time:.2f}秒 后触发")
            # 创建定时器，在等待时间结束后触发上货操作
            load_timer = threading.Timer(wait_time, trigger_load, args=(mqtt_client, retry_count, max_retries))
            load_timer.start()
            return

        try:
            # 尝试将上货推杆对象放入队列中，如果队列已满则立即抛出异常而不阻塞
            load_rod_queue.put_nowait(load_rod)
            # 记录本次上货操作的时间戳
            last_load_time = time.time()
            # 记录上货推杆触发日志信息
            logger.info("触发上货推杆")

        except queue.Full:
            # 处理上料队列满的情况，采用指数退避策略进行重试
            # 当重试次数未达到上限时，计算等待时间并启动定时重试机制
            # 若重试次数已达上限，则记录错误日志并停止上货操作
            if retry_count < max_retries:
                # 计算指数退避等待时间，最小值为当前重试次数对应的2的幂次方秒，最大不超过5秒
                wait_time = min(1.0 * (2 ** retry_count), 5.0)  # 指数退避，最大5秒
                logger.warning(f"上料队列满，{wait_time:.2f}秒后重试上货（第 {retry_count + 1}/{max_retries} 次）")
                # 创建并启动定时器，在指定等待时间后重新触发上货操作
                load_timer = threading.Timer(wait_time, trigger_load, args=(mqtt_client, retry_count + 1, max_retries))
                load_timer.start()
            else:
                logger.error("上料队列满，重试达到上限，暂停上货")


if __name__ == '__main__':
    # 初始化 MQTT 客户端，连接到指定的 IP 和端口，并使用用户名和密码进行认证
    mqtt_client = Mqtt_Clt("127.0.0.1", 21883, "bb", "aa", 60)
    logger.info("系统启动，初始化设备")

    # 控制传送带运行和推杆全部拉回初始状态
    mqtt_client.control_device("conveyor", 'run')
    mqtt_client.control_device("rod_control", "all_pull")

    # 初始化货物追踪器和工作线程
    tracker = CargoTracker()
    workers = start_workers(mqtt_client, tracker)

    # 启动定期清理线程：每30秒调用一次 tracker.cleanup_old_cargos() 清理过期货物信息
    def periodic_cleanup():
        while True:
            tracker.cleanup_old_cargos()
            time.sleep(30)
    cleanup_thread = threading.Thread(target=periodic_cleanup, daemon=True)
    cleanup_thread.start()

    # 执行首次上货操作
    logger.info("系统启动，执行首次上货")
    trigger_load(mqtt_client)

    try:
        while True:
            try:
                # 从 MQTT 消息队列中获取消息
                msg = mqtt_client.mqtt_queue.get()
            except Exception as e:
                logger.error(f"MQTT 队列获取失败: {e}")
                continue

            # 判断是否为图像消息
            if 'image' in msg:
                try:
                    logger.info("收到图像，开始解码")
                    # 解码 base64 编码的图像数据
                    image_data = base64.b64decode(msg['image'])
                    image_array = np.frombuffer(image_data, np.uint8)
                    image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
                    if image is None:
                        logger.error("图像解码失败")
                        continue

                    try:
                        # 使用 YOLO 模型对图像进行分类预测
                        classification = predict_yolo(image)
                    except Exception as e:
                        logger.error(f"YOLO 推理失败: {e}")
                        continue

                    logger.info(f"识别结果: {classification}")
                    # 如果识别结果在预定义分类映射中，则创建新货物并触发上货
                    if classification in CLASS_SWITCH_MAP:
                        cargo_id = tracker.create_new_cargo(classification)
                        trigger_load(mqtt_client)
                    else:
                        logger.warning(f"忽略无效分类: {classification}")
                except Exception as e:
                    logger.error(f"图像处理错误: {e}")
                    continue

            else:
                logger.info(f"收到传感器消息: {msg}")
                # 提取有效的传感器触发信号
                valid_switches = [s for s in CLASS_SWITCH_MAP.values() if msg.get(s, False)]
                if not valid_switches:
                    logger.warning(f"无有效传感器触发，忽略消息: {msg}")
                    continue

                # 按照预设顺序对有效传感器进行排序
                sorted_valid_switches = sorted(valid_switches, key=lambda s: SENSOR_ORDER.index(s))
                if len(valid_switches) > 1:
                    logger.warning(f"异常: 多个传感器同时触发 {valid_switches}，按顺序处理: {sorted_valid_switches}")
                else:
                    logger.info(f"有效传感器: {sorted_valid_switches[0]}（硬件True）")

                # 依次处理每个有效传感器触发事件
                for active_switch in sorted_valid_switches:
                    # 更新货物状态并获取对应的货物 ID
                    cargo_id = tracker.update_cargo_switch(active_switch)
                    if cargo_id:
                        # 根据货物分类获取对应的推杆名称
                        rod_to_push = CLASS_ROD_MAP.get(tracker.cargo_map[cargo_id]["classification"])
                        if rod_to_push:
                            logger.info(f"货物ID: {cargo_id} 触发分类推杆: {rod_to_push}")
                            task = (cargo_id, rod_to_push, active_switch)
                            try:
                                # 将任务放入分类推杆队列
                                other_rod_queue.put_nowait(task)
                            except queue.Full:
                                logger.warning("分类推杆队列满，1秒后重试放入任务")
                                threading.Timer(1.0, lambda: other_rod_queue.put_nowait(task)).start()
                        else:
                            logger.error(f"货物ID: {cargo_id} 无对应推杆")

    except KeyboardInterrupt:
        logger.info("收到终止信号，停止系统")
        if load_timer is not None:
            load_timer.cancel()
    finally:
        # 停止所有工作线程并清理资源
         # 向队列中添加结束信号，通知加载杆处理线程退出
        load_rod_queue.put(None)
        # 为其他杆处理线程添加结束信号
        for _ in range(len(workers) - 1):
            other_rod_queue.put(None)
        # 等待所有工作线程完成并退出
        for w in workers:
            w.join()
        # 停止传送带设备
        mqtt_client.control_device("conveyor", 'stop')
        # 控制所有杆设备进行收回操作
        mqtt_client.control_device("rod_control", "all_pull")
        # 清空货物映射数据
        tracker.cargo_map.clear()
        logger.info("系统已退出")

